import argparse
import os
import sys
from pathlib import Path

import torch
from mmengine.runner import Runner
from mmengine.optim import OptimWrapper, OptimWrapperDict

from gan.mnist import MNISTDataset
from gan.model import create_gan

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

NUM_WORKERS = int(os.cpu_count() / 2)
BATCH_SIZE = 256 if torch.cuda.is_available() else 64


def main(opt):
    # train dataloader
    dataset = MNISTDataset(opt.data_dir, [])
    train_dataloader = dict(
        batch_size=BATCH_SIZE,
        num_workers=NUM_WORKERS,
        persistent_workers=True,
        sampler=dict(type="DefaultSampler", shuffle=True),
        dataset=dataset,
    )
    train_dataloader = Runner.build_dataloader(train_dataloader)

    # GAN
    model, generator, discriminator = create_gan(None)

    # Multi optimizers
    opt_g = torch.optim.Adam(generator.parameters(), lr=0.0001, betas=(0.5, 0.999))
    opt_g_wrapper = OptimWrapper(opt_g)

    opt_d = torch.optim.Adam(discriminator.parameters(), lr=0.0001, betas=(0.5, 0.999))
    opt_d_wrapper = OptimWrapper(opt_d)

    opt_wrapper_dict = OptimWrapperDict(
        generator=opt_g_wrapper, discriminator=opt_d_wrapper
    )

    # Train
    train_cfg = dict(by_epoch=True, max_epochs=opt.epochs)
    runner = Runner(
        model,
        work_dir=opt.work_dir,
        train_dataloader=train_dataloader,
        train_cfg=train_cfg,
        optim_wrapper=opt_wrapper_dict,
    )
    runner.train()


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", default=ROOT / "data", help="the data directory")
    parser.add_argument(
        "--work_dir", default=ROOT / "runs/gan", help="the working directory"
    )
    parser.add_argument("--epochs", default=50, type=int, help="epochs to train")
    opt = parser.parse_args()
    return opt


if __name__ == "__main__":
    opt = parse_opt()
    main(opt)
